import argparse
import json
import math
import os, cv2

import numpy as np
from tqdm import tqdm


parser = argparse.ArgumentParser('rec')
parser.add_argument("--train_file_path", type=str, default='/mnt/train_data/vin/Steel_Seal/test.txt')
parser.add_argument("--imgpath", type=str, default='/mnt/train_data/vin/Steel_Seal/test_img')
parser.add_argument("--savepath", type=str, default='/mnt/train_data/vin/Steel_Seal/test_cut')

args = parser.parse_args()

if not os.path.exists(args.savepath):
    os.mkdir(args.savepath)

def distance(x,y):
    return  math.sqrt((x[0]-y[0])**2+(x[1]-y[1])**2)

def get_center(p1,p2,p3,p4):
    cx = int((p1[0]+p2[0]+p3[0]+p4[0])/4.0)
    cy = int((p1[1]+p2[1]+p3[1]+p4[1])/4.0)
    return [cx, cy]

def rotate(ps, m):
    pts = np.float32(ps).reshape([-1, 2])  # 要映射的点
    pts = np.hstack([pts, np.ones([len(pts), 1])]).T
    target_point = np.dot(m, pts)
    target_point = [[int(target_point[0][x]),int(target_point[1][x])] for x in range(len(target_point[0]))]
    return target_point

def rotate_img_and_point(img,points,angle,center_x,center_y,resize_rate=1.0):
    h,w,c = img.shape
    M = cv2.getRotationMatrix2D((center_x,center_y), angle, resize_rate)
    res_img = cv2.warpAffine(img, M, (w, h))
    out_points = rotate(points, M)
    return res_img,out_points

def rotate_point(points,angle,center_x,center_y,resize_rate=1.0):
    M = cv2.getRotationMatrix2D((center_x, center_y), angle, resize_rate)
    out_points = rotate(points, M)
    return  out_points
def process_(img, sorts):


    sp = img.shape
    if distance(sorts[0], sorts[1]) > distance(sorts[1], sorts[2]):
        w = distance(sorts[0], sorts[1])
        h = distance(sorts[1], sorts[2])
        if sorts[1][1] > sorts[0][1]:
            angel = math.acos((sorts[1][0] - sorts[0][0]) / distance(sorts[0], sorts[1]))
            # rotate_img_and_point(img,points,angle,center_x,center_y,resize_rate=1.0):
            src_img, p = rotate_img_and_point(img, get_center(sorts[0], sorts[1], sorts[2], sorts[3]),
                                              (angel / math.pi * 180), sp[1] / 2.0, sp[0] / 2.0)


        else:

            if sorts[1][0] < sorts[0][0]:
                angel = math.acos((sorts[1][0] - sorts[0][0]) / distance(sorts[0], sorts[1]))
                src_img, p = rotate_img_and_point(img, get_center(sorts[0], sorts[1], sorts[2], sorts[3]),
                                                  (-1) * (angel / math.pi * 180), sp[1] / 2.0, sp[0] / 2.0)


            else:

                angel = math.acos((sorts[1][0] - sorts[0][0]) / distance(sorts[0], sorts[1]))
                src_img, p = rotate_img_and_point(img, get_center(sorts[0], sorts[1], sorts[2], sorts[3]),
                                                  (-1) * (angel / math.pi * 180), sp[1] / 2.0, sp[0] / 2.0)


    else:
        h = distance(sorts[0], sorts[1])
        w = distance(sorts[1], sorts[2])
        if sorts[1][0] > sorts[2][0]:
            angel = math.acos((sorts[1][0] - sorts[2][0]) / distance(sorts[1], sorts[2]))

            src_img, p = rotate_img_and_point(img, get_center(sorts[0], sorts[1], sorts[2], sorts[3]),
                                              (-1) * (angel / math.pi * 180), sp[1] / 2.0, sp[0] / 2.0)


        else:
            angel = math.acos((sorts[2][0] - sorts[1][0]) / distance(sorts[1], sorts[2]))
            src_img, p = rotate_img_and_point(img, get_center(sorts[0], sorts[1], sorts[2], sorts[3]),
                                              ((angel / math.pi) * 180), sp[1] / 2.0, sp[0] / 2.0)


    result_p = []
    res_sp = src_img.shape
    if int((p[0][1] - int(h / 2.0)) ) < 0:
        result_p.append(0)
    else:
        result_p.append(int((p[0][1] - int(h / 2.0)) ))

    if int((p[0][1] + int(h / 2.0)) ) > res_sp[0]:
        result_p.append(res_sp[0])
    else:
        result_p.append(int((p[0][1] + int(h / 2.0)) ))

    if int((p[0][0] - int(w / 2.0)) ) < 0:
        result_p.append(0)
    else:
        result_p.append(int((p[0][0] - int(w / 2.0)) ))

    if int((p[0][0] + int(w / 2.0))) > res_sp[1]:
        result_p.append(res_sp[1])
    else:
        result_p.append(int((p[0][0] + int(w / 2.0)) ))

    src_img = src_img[result_p[0]:result_p[1], result_p[2]:result_p[3]]
    return src_img



def paixu(p):
    p = sorted(p, key=lambda point: point[1])
    sorts = []
    if p[0][0]<p[1][0]:
        sorts.append(p[0])
        sorts.append(p[1])
    else:
        sorts.append(p[1])
        sorts.append(p[0])
    if p[2][0]<p[3][0]:
        sorts.append(p[3])
        sorts.append(p[2])
    else:
        sorts.append(p[2])
        sorts.append(p[3])
    return sorts


img_list = []

with open(args.train_file_path,"r",encoding='utf-8') as f:
    lines = f.readlines()
    for i in tqdm(lines):
        img_path, json_txt = i.split("\t")
        json_ = json.loads(json_txt)
        img_name = os.path.basename(img_path)
        name = img_name.split(".")[0]
        for j in (range(len(json_))):
            img =cv2.imread(os.path.join(args.imgpath, img_name))
            points = paixu(json_[j]['points'])
            src_img = process_(img,points)
            cv2.imwrite(os.path.join(args.savepath,name+"_"+str(j)+".jpg"),src_img)





